from python_ai.common.xcommon import sep
import warnings
import numpy as np
import matplotlib.pyplot as plt
import tensorflow.compat.v1 as tf

warnings.simplefilter('error', UserWarning)

np.random.seed(1)
tf.set_random_seed(777)
plt.figure(figsize=[16, 8])
spr = 4  # subplot row
spc = 8  # subplot column
spn = 0

# load data
path = r'../../../../../large_data/ML1/hand_writing/'
data = np.loadtxt(path + 'imgX.txt', delimiter=',')
y = np.loadtxt(path + 'labely.txt', delimiter=',')
m, n = data.shape
nn = int(np.sqrt(n))
print(m, n)
print(f'{nn} x {nn}')
y[y == 10] = 0  # ATTENTION

# shuffle
rand_idx = np.random.permutation(m)
data = data[rand_idx]
y = y[rand_idx]

# placeholder
input = tf.placeholder(tf.float32, [None, nn, nn, 1], 'placeholder_input')

# convolution kernels
knl_raw = np.array([
    [[1, 0, -1],
     [1, 0, -1],
     [1, 0, -1],],
    [[1, 1, 1],
     [0, 0, 0],
     [-1, -1, -1]],
    [[-1, 0, 1],
     [-1, 0, 1],
     [-1, 0, 1],],
    [[-1, -1, -1],
     [0, 0, 0],
     [1, 1, 1]],
], dtype=np.float32)
print(f'Kernel data shape: {knl_raw.shape}')
kh, kw, cin, cout = knl_raw.shape[1], knl_raw.shape[2], 1, knl_raw.shape[0]
print(f'Target shape: {kh, kw, cin, cout}')
knl = np.zeros([kh, kw, cin, cout])
for c_idx, c in enumerate(knl_raw):  # ATTENTION Transpose and expand_dims is much better than this.
    for h_idx, h in enumerate(c):
        for w_idx, w in enumerate(h):
            knl[h_idx, w_idx, 0, c_idx] = w
filter1 = tf.Variable(knl, dtype=tf.float32, name='filter1')
print(filter1)

# ATTENTION
# input: `"NHWC" or "NCHW"`
# filter: A 4-D tensor of shape
#       `[filter_height, filter_width, in_channels, out_channels]`
# padding:
#       VALID: No padding
#           20x20=strides1x1=>18x18
#           20x20=strides2x2=>9x9
#           7x7=strides1x1=>5x5
#           7x7=strides2x2=>3x3
#       SAME: Padding to the same size of input
#           20x20=strides1x1=>20x20
#           20x20=strides2x2=>10x10
#           7x7=strides1x1=>7x7
#           7x7=strides2x2=>4x4
# strides: An int or list of `ints` that has length `1`, `2` or `4`.  The
#       stride of the sliding window for each dimension of `input`. If a single
#       value is given it is replicated in the `H` and `W` dimension. By default
#       the `N` and `C` dimensions are set to 1. The dimension order is determined
#       by the value of `data_format`, see below for details.
# Returns:
#     A `Tensor`. Has the same type as `input`.
result_1_valid = tf.nn.conv2d(input, filter1, strides=[1, 1, 1, 1], padding='VALID')
print(result_1_valid)
result_1_same = tf.nn.conv2d(input, filter1, strides=[1, 1, 1, 1], padding='SAME')
print(result_1_same)
result_2_valid = tf.nn.conv2d(input, filter1, strides=[1, 2, 2, 1], padding='VALID')
print(result_2_valid)
result_2_same = tf.nn.conv2d(input, filter1, strides=[1, 2, 2, 1], padding='SAME')
print(result_2_same)
result = [
    result_1_valid,
    # result_1_same,
    # result_2_valid,
    # result_2_same
]
print(result)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())

    # show detail
    def output(tgt, num_all):
        global spn
        num = 0
        idx = -1
        while True:
            idx += 1
            if y[idx] != tgt:
                continue
            num += 1
            if num > num_all:
                break
            sep((tgt, num))
            spn += 1
            if spn > spr * spc:
                break
            plt.subplot(spr, spc, spn)
            plt.title(str(y[idx]))
            pic = data[idx].reshape(nn, nn).T
            pic_mat = pic.reshape(1, nn, nn, 1)
            plt.imshow(pic)

            pics_list = sess.run(result, feed_dict={input: pic_mat})
            for m_pic in pics_list:
                for pics in m_pic:
                    h, w, conv = pics.shape
                    print(pics.shape)
                    for i in range(conv):
                        spn += 1
                        if spn > spr * spc:
                            break
                        plt.subplot(spr, spc, spn)
                        p = pics[:, :, i]
                        plt.imshow(p)


    # nums = [0, 1, 5, 7]
    nums = range(10)
    for i in nums:
        output(i, 1)

    plt.show()
